from utils import paddle_aux
import os
import paddle
"""
Validate a trained YOLOv5 model accuracy on a custom dataset

Usage:
    $ python path/to/val.py --weights yolov5s.pt --data coco128.yaml --img 640

Usage - formats:
    $ python path/to/val.py --weights yolov5s.pt                 # PyTorch
                                      yolov5s.torchscript        # TorchScript
                                      yolov5s.onnx               # ONNX Runtime or OpenCV DNN with --dnn
                                      yolov5s.xml                # OpenVINO
                                      yolov5s.engine             # TensorRT
                                      yolov5s.mlmodel            # CoreML (MacOS-only)
                                      yolov5s_saved_model        # TensorFlow SavedModel
                                      yolov5s.pb                 # TensorFlow GraphDef
                                      yolov5s.tflite             # TensorFlow Lite
                                      yolov5s_edgetpu.tflite     # TensorFlow Edge TPU
"""
import argparse
import json
import sys
from pathlib import Path
from threading import Thread
import numpy as np
from tqdm import tqdm
FILE = Path(__file__).resolve()
ROOT = FILE.parents[0]
if str(ROOT) not in sys.path:
    sys.path.append(str(ROOT))
ROOT = Path(os.path.relpath(ROOT, Path.cwd()))
from models.common import DetectMultiBackend
from utils.callbacks import Callbacks
from utils.datasets import create_dataloader
from utils.general import LOGGER, box_iou, check_dataset, check_img_size, check_requirements, check_yaml, coco80_to_coco91_class, colorstr, increment_path, non_max_suppression, print_args, scale_coords, xywh2xyxy, xyxy2xywh
from utils.metrics import ConfusionMatrix, ap_per_class
from utils.plots import output_to_target, plot_images, plot_val_study
from utils.torch_utils import select_device, time_sync


def save_one_txt(predn, save_conf, shape, file):
    gn = paddle.to_tensor(data=shape)[[1, 0, 1, 0]]
    for *xyxy, conf, cls in predn.tolist():
        xywh = (xyxy2xywh(paddle.to_tensor(data=xyxy).view(1, 4)) / gn).view(-1
            ).tolist()
        line = (cls, *xywh, conf) if save_conf else (cls, *xywh)
        with open(file, 'a') as f:
            f.write(('%g ' * len(line)).rstrip() % line + '\n')


def save_one_json(predn, jdict, path, class_map):
    image_id = int(path.stem) if path.stem.isnumeric() else path.stem
    box = xyxy2xywh(predn[:, :4])
    box[:, :2] -= box[:, 2:] / 2
    for p, b in zip(predn.tolist(), box.tolist()):
        jdict.append({'image_id': image_id, 'category_id': class_map[int(p[
            5])], 'bbox': [round(x, 3) for x in b], 'score': round(p[4], 5)})


def process_batch(detections, labels, iouv):
    """
    Return correct predictions matrix. Both sets of boxes are in (x1, y1, x2, y2) format.
    Arguments:
        detections (Array[N, 6]), x1, y1, x2, y2, conf, class
        labels (Array[M, 5]), class, x1, y1, x2, y2
    Returns:
        correct (Array[N, 10]), for 10 IoU levels
    """
    correct = paddle.zeros(shape=[tuple(detections.shape)[0], tuple(iouv.
        shape)[0]], dtype='bool')
    iou = box_iou(labels[:, 1:], detections[:, :4])
    x = paddle.where((iou >= iouv[0]) & (labels[:, 0:1] == detections[:, 5]))
    if tuple(x[0].shape)[0]:
        matches = paddle.concat(x=(paddle.stack(x=x, axis=1), iou[x[0], x[1
            ]][:, None]), axis=1).cpu().numpy()
        if tuple(x[0].shape)[0] > 1:
            matches = matches[matches[:, 2].argsort()[::-1]]
            matches = matches[np.unique(matches[:, 1], return_index=True)[1]]
            matches = matches[np.unique(matches[:, 0], return_index=True)[1]]
        matches = paddle.to_tensor(data=matches).to(iouv.place)
        correct[matches[:, 1].astype(dtype='int64')] = matches[:, 2:3] >= iouv
    return correct


@paddle.no_grad()
def run(data, weights=None, batch_size=32, imgsz=640, conf_thres=0.001,
    iou_thres=0.6, task='val', device='', workers=8, single_cls=False,
    augment=False, verbose=False, save_txt=False, save_hybrid=False,
    save_conf=False, save_json=False, project=ROOT / 'runs/val', name='exp',
    exist_ok=False, half=True, dnn=False, model=None, dataloader=None,
    save_dir=Path(''), plots=True, callbacks=Callbacks(), compute_loss=None):
    training = model is not None
    if training:
        device, pt, jit, engine = next(model.parameters()
            ).place, True, False, False
        half &= device.type != 'cpu'
        model.astype(dtype='float16') if half else model.astype(dtype='float32'
            )
    else:
        device = select_device(device, batch_size=batch_size)
        save_dir = increment_path(Path(project) / name, exist_ok=exist_ok)
        (save_dir / 'labels' if save_txt else save_dir).mkdir(parents=True,
            exist_ok=True)
        model = DetectMultiBackend(weights, device=device, dnn=dnn, data=data)
        stride, pt, jit, onnx, engine = (model.stride, model.pt, model.jit,
            model.onnx, model.engine)
        imgsz = check_img_size(imgsz, s=stride)
        half &= (pt or jit or onnx or engine) and device.type != 'cpu'
        if pt or jit:
            model.model.half() if half else model.model.float()
        elif engine:
            batch_size = model.batch_size
            if model.trt_fp16_input != half:
                LOGGER.info('model ' + ('requires' if model.trt_fp16_input else
                    'incompatible with') + ' --half. Adjusting automatically.')
                half = model.trt_fp16_input
        else:
            half = False
            batch_size = 1
            device = paddle.CPUPlace()
            LOGGER.info(
                f'Forcing --batch-size 1 square inference shape(1,3,{imgsz},{imgsz}) for non-PyTorch backends'
                )
        data = check_dataset(data)
    model.eval()
    is_coco = isinstance(data.get('val'), str) and data['val'].endswith(
        'coco/val2017.txt')
    nc = 1 if single_cls else int(data['nc'])
    iouv = paddle.linspace(start=0.5, stop=0.95, num=10).to(device)
    niou = iouv.size
    if not training:
        model.warmup(imgsz=(1 if pt else batch_size, 3, imgsz, imgsz), half
            =half)
        pad = 0.0 if task in ('speed', 'benchmark') else 0.5
        rect = False if task == 'benchmark' else pt
        task = task if task in ('train', 'val', 'test') else 'val'
        dataloader = create_dataloader(data[task], imgsz, batch_size,
            stride, single_cls, pad=pad, rect=rect, workers=workers, prefix
            =colorstr(f'{task}: '))[0]
    seen = 0
    confusion_matrix = ConfusionMatrix(nc=nc)
    """Class Attribute: torch.Tensor.names, can not convert, please check whether it is torch.Tensor.*/torch.autograd.function.FunctionCtx.*/torch.distributions.Distribution.* and convert manually"""
    names = {k: v for k, v in enumerate(model.names if hasattr(model,
        'names') else model.module.names)}
    class_map = coco80_to_coco91_class() if is_coco else list(range(1000))
    s = ('%20s' + '%11s' * 6) % ('Class', 'Images', 'Labels', 'P', 'R',
        'mAP@.5', 'mAP@.5:.95')
    dt, p, r, f1, mp, mr, map50, map = [0.0, 0.0, 0.0
        ], 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0
    loss = paddle.zeros(shape=[3])
    jdict, stats, ap, ap_class = [], [], [], []
    pbar = tqdm(dataloader, desc=s, bar_format=
        '{l_bar}{bar:10}{r_bar}{bar:-10b}')
    for batch_i, (im, targets, paths, shapes) in enumerate(pbar):
        t1 = time_sync()
        if pt or jit or engine:
            im = im.to(device, blocking=not True)
            targets = targets.to(device)
        im = im.astype(dtype='float16') if half else im.astype(dtype='float32')
        im /= 255
        nb, _, height, width = tuple(im.shape)
        t2 = time_sync()
        dt[0] += t2 - t1
        out, train_out = model(im) if training else model(im, augment=
            augment, val=True)
        dt[1] += time_sync() - t2
        if compute_loss:
            loss += compute_loss([x.astype(dtype='float32') for x in
                train_out], targets)[1]
        targets[:, 2:] *= paddle.to_tensor(data=[width, height, width,
            height], dtype='float32').to(device)
        lb = [targets[targets[:, 0] == i, 1:] for i in range(nb)
            ] if save_hybrid else []
        t3 = time_sync()
        out = non_max_suppression(out, conf_thres, iou_thres, labels=lb,
            multi_label=True, agnostic=single_cls)
        dt[2] += time_sync() - t3
        for si, pred in enumerate(out):
            labels = targets[targets[:, 0] == si, 1:]
            nl = len(labels)
            tcls = labels[:, 0].tolist() if nl else []
            path, shape = Path(paths[si]), shapes[si][0]
            seen += 1
            if len(pred) == 0:
                if nl:
                    stats.append((paddle.zeros(shape=[0, niou], dtype=
                        'bool'), paddle.to_tensor(data=[]), paddle.
                        to_tensor(data=[]), tcls))
                continue
            if single_cls:
                pred[:, 5] = 0
            predn = pred.clone()
            scale_coords(tuple(im[si].shape)[1:], predn[:, :4], shape,
                shapes[si][1])
            if nl:
                tbox = xywh2xyxy(labels[:, 1:5])
                scale_coords(tuple(im[si].shape)[1:], tbox, shape, shapes[
                    si][1])
                labelsn = paddle.concat(x=(labels[:, 0:1], tbox), axis=1)
                correct = process_batch(predn, labelsn, iouv)
                if plots:
                    confusion_matrix.process_batch(predn, labelsn)
            else:
                correct = paddle.zeros(shape=[tuple(pred.shape)[0], niou],
                    dtype='bool')
            stats.append((correct.cpu(), pred[:, 4].cpu(), pred[:, 5].cpu(),
                tcls))
            if save_txt:
                save_one_txt(predn, save_conf, shape, file=save_dir /
                    'labels' / (path.stem + '.txt'))
            if save_json:
                save_one_json(predn, jdict, path, class_map)
            callbacks.run('on_val_image_end', pred, predn, path, names, im[si])
        if plots and batch_i < 3:
            f = save_dir / f'val_batch{batch_i}_labels.jpg'
            Thread(target=plot_images, args=(im, targets, paths, f, names),
                daemon=True).start()
            f = save_dir / f'val_batch{batch_i}_pred.jpg'
            Thread(target=plot_images, args=(im, output_to_target(out),
                paths, f, names), daemon=True).start()
    stats = [np.concatenate(x, 0) for x in zip(*stats)]
    if len(stats) and stats[0].astype('bool').any():
        tp, fp, p, r, f1, ap, ap_class = ap_per_class(*stats, plot=plots,
            save_dir=save_dir, names=names)
        ap50, ap = ap[:, 0], ap.mean(axis=1)
        mp, mr, map50, map = p.mean(), r.mean(), ap50.mean(), ap.mean()
        nt = np.bincount(stats[3].astype(np.int64), minlength=nc)
    else:
        nt = paddle.zeros(shape=[1])
    pf = '%20s' + '%11i' * 2 + '%11.3g' * 4
    LOGGER.info(pf % ('all', seen, nt.sum(), mp, mr, map50, map))
    if (verbose or nc < 50 and not training) and nc > 1 and len(stats):
        for i, c in enumerate(ap_class):
            LOGGER.info(pf % (names[c], seen, nt[c], p[i], r[i], ap50[i],
                ap[i]))
    t = tuple(x / seen * 1000.0 for x in dt)
    if not training:
        shape = batch_size, 3, imgsz, imgsz
        LOGGER.info(
            f'Speed: %.1fms pre-process, %.1fms inference, %.1fms NMS per image at shape {shape}'
             % t)
    if plots:
        confusion_matrix.plot(save_dir=save_dir, names=list(names.values()))
        callbacks.run('on_val_end')
    if save_json and len(jdict):
        w = Path(weights[0] if isinstance(weights, list) else weights
            ).stem if weights is not None else ''
        anno_json = str(Path(data.get('path', '../coco')) /
            'annotations/instances_val2017.json')
        pred_json = str(save_dir / f'{w}_predictions.json')
        LOGGER.info(f'\nEvaluating pycocotools mAP... saving {pred_json}...')
        with open(pred_json, 'w') as f:
            json.dump(jdict, f)
        try:
            check_requirements(['pycocotools'])
            from pycocotools.coco import COCO
            from pycocotools.cocoeval import COCOeval
            anno = COCO(anno_json)
            pred = anno.loadRes(pred_json)
            eval = COCOeval(anno, pred, 'bbox')
            if is_coco:
                eval.params.imgIds = [int(Path(x).stem) for x in dataloader
                    .dataset.im_files]
            eval.evaluate()
            eval.accumulate()
            eval.summarize()
            map, map50 = eval.stats[:2]
        except Exception as e:
            LOGGER.info(f'pycocotools unable to run: {e}')
    model.astype(dtype='float32')
    if not training:
        s = (
            f"\n{len(list(save_dir.glob('labels/*.txt')))} labels saved to {save_dir / 'labels'}"
             if save_txt else '')
        LOGGER.info(f"Results saved to {colorstr('bold', save_dir)}{s}")
    maps = np.zeros(nc) + map
    for i, c in enumerate(ap_class):
        maps[c] = ap[i]
    return (mp, mr, map50, map, *(loss.cpu() / len(dataloader)).tolist()
        ), maps, t


def parse_opt():
    parser = argparse.ArgumentParser()
    parser.add_argument('--data', type=str, default=ROOT / 'data/VOC.yaml',
        help='dataset.yaml path')
    parser.add_argument('--weights', nargs='+', type=str, default=ROOT /
        'runs/train/yolov5s-mobilenetV2improved-jingshan20220816/weights/best.pt'
        , help='model.pt path(s)')
    parser.add_argument('--batch-size', type=int, default=8, help='batch size')
    parser.add_argument('--imgsz', '--img', '--img-size', type=int, default
        =1024, help='inference size (pixels)')
    parser.add_argument('--conf-thres', type=float, default=0.6, help=
        'confidence threshold')
    parser.add_argument('--iou-thres', type=float, default=0.6, help=
        'NMS IoU threshold')
    parser.add_argument('--task', default='test', help=
        'train, val, test, speed or study')
    parser.add_argument('--device', default='0', help=
        'cuda device, i.e. 0 or 0,1,2,3 or cpu')
    parser.add_argument('--workers', type=int, default=8, help=
        'max dataloader workers (per RANK in DDP mode)')
    parser.add_argument('--single-cls', action='store_true', help=
        'treat as single-class dataset')
    parser.add_argument('--augment', action='store_true', help=
        'augmented inference')
    parser.add_argument('--verbose', action='store_true', help=
        'report mAP by class')
    parser.add_argument('--save-txt', action='store_true', help=
        'save results to *.txt')
    parser.add_argument('--save-hybrid', action='store_true', help=
        'save label+prediction hybrid results to *.txt')
    parser.add_argument('--save-conf', action='store_true', help=
        'save confidences in --save-txt labels')
    parser.add_argument('--save-json', action='store_true', help=
        'save a COCO-JSON results file')
    parser.add_argument('--project', default=ROOT / 'runs/val', help=
        'save to project/name')
    parser.add_argument('--name', default=
        'yolov5s-mobilenetV2improved-jingshan20220816', help=
        'save to project/name')
    parser.add_argument('--exist-ok', action='store_true', help=
        'existing project/name ok, do not increment')
    parser.add_argument('--half', action='store_true', help=
        'use FP16 half-precision inference')
    parser.add_argument('--dnn', action='store_true', help=
        'use OpenCV DNN for ONNX inference')
    opt = parser.parse_args()
    opt.data = check_yaml(opt.data)
    opt.save_json |= opt.data.endswith('coco.yaml')
    opt.save_txt |= opt.save_hybrid
    print_args(FILE.stem, opt)
    return opt


def main(opt):
    check_requirements(requirements=ROOT / 'requirements.txt', exclude=(
        'tensorboard', 'thop'))
    if opt.task in ('train', 'val', 'test'):
        if opt.conf_thres > 0.001:
            LOGGER.info(
                f'WARNING: confidence threshold {opt.conf_thres} >> 0.001 will produce invalid mAP values.'
                )
        run(**vars(opt))
    else:
        weights = opt.weights if isinstance(opt.weights, list) else [opt.
            weights]
        opt.half = True
        if opt.task == 'speed':
            opt.conf_thres, opt.iou_thres, opt.save_json = 0.25, 0.45, False
            for opt.weights in weights:
                run(**vars(opt), plots=False)
        elif opt.task == 'study':
            for opt.weights in weights:
                f = f'study_{Path(opt.data).stem}_{Path(opt.weights).stem}.txt'
                x, y = list(range(256, 1536 + 128, 128)), []
                for opt.imgsz in x:
                    LOGGER.info(f'\nRunning {f} --imgsz {opt.imgsz}...')
                    r, _, t = run(**vars(opt), plots=False)
                    y.append(r + t)
                np.savetxt(f, y, fmt='%10.4g')
            os.system('zip -r study.zip study_*.txt')
            plot_val_study(x=x)


if __name__ == '__main__':
    opt = parse_opt()
    main(opt)
